#include "rbtree.hpp"
#include <utility> // for std::pair

namespace STL
{
    template <class Key, class Val>
    class KeyOfT
    {
    public:
        Key operator()(const std::pair<Key, Val> &val) const
        {
            return val.first;
        }
    };

    template <class Key, class Val>
    class map
    {
    private:
        using ValueType = std::pair<Key, Val>;
        using BaseRBTree = RBTree<Key, ValueType, KeyOfT<Key, Val>, std::less<Key>>;

    public:
        using iterator = typename BaseRBTree::iterator;

        iterator begin()
        {
            return _rbtree.begin();
        }

        iterator end()
        {
            return _rbtree.end();
        }

        bool insert(const ValueType &val)
        {
            return _rbtree.insert(val);
        }

        // 添加一些常用方法
        bool insert(const Key &key, const Val &value)
        {
            return _rbtree.insert(std::make_pair(key, value));
        }

        Val &operator[](const Key &key)
        {
            // 简化版本，实际应该先查找再插入
            auto result = _rbtree.insert(std::make_pair(key, Val()));
            return (*result.first).second; // 返回值的引用
        }

        size_t size() const
        {
            // 需要在实际的RBTree中实现size()
            size_t count = 0;
            auto it = _rbtree.begin();
            while (it != _rbtree.end())
            {
                ++count;
                ++it;
            }
            return count;
        }

        bool empty() const
        {
            return _rbtree.begin() == _rbtree.end();
        }

        // 添加 inorder 方法
        void inorder() const { _rbtree.inorder(); }

    private:
        BaseRBTree _rbtree;
    };
}